﻿// UDPReceiverSample.cpp: 定义应用程序的入口点。
//

#include "udpreceiver.h"
#include <iomanip>
#include <cstring>

#ifdef _WIN32
#include <winsock2.h>
#include <mswsock.h>
#include <windows.h>
#else
#include <arpa/inet.h>
#endif

using namespace std;

#ifdef _WIN32
// 静态实例指针定义
UdpReceiver* UdpReceiver::instance_ = nullptr;
#endif

// UdpReceiver 类实现
UdpReceiver::UdpReceiver(asio::io_context& io_context, unsigned short port)
    : io_context_(io_context)
    , socket_(io_context)
    , signals_(io_context)
    , statistics_timer_(io_context)
    , start_time_(std::chrono::steady_clock::now())
    , last_statistics_time_(std::chrono::steady_clock::now()) {
    
#ifdef _WIN32
    // 设置静态实例指针
    instance_ = this;
    
    // 设置Windows Console Control Handler
    if (!SetConsoleCtrlHandler(console_ctrl_handler, TRUE)) {
        std::cerr << "警告: 无法设置Console Control Handler" << std::endl;
    }
#endif
    
    try {
        // 打开socket
        socket_.open(udp::v4());
        
        // 设置socket选项以优化性能，减少丢包
        socket_.set_option(asio::socket_base::reuse_address(true)); // 允许端口重用
        socket_.set_option(asio::socket_base::receive_buffer_size(UDP_SYSTEM_RCV_BUFFER)); // 接收缓冲区
        
        // 绑定端口
        socket_.bind(udp::endpoint(udp::v4(), port));
        
#ifdef _WIN32
        // Windows特定优化
        SOCKET native_socket = socket_.native_handle();
        
        // 设置更大的接收缓冲区（UDP专用）
        int buffer_size = UDP_SYSTEM_RCV_BUFFER;
        setsockopt(native_socket, SOL_SOCKET, SO_RCVBUF, (char*)&buffer_size, sizeof(buffer_size));
        
        // 设置发送缓冲区
        int send_buffer_size = UDP_SYSTEM_SND_BUFFER;
        setsockopt(native_socket, SOL_SOCKET, SO_SNDBUF, (char*)&send_buffer_size, sizeof(send_buffer_size));
        
        std::cout << "UDP 服务器启动 (Windows优化)，监听端口: " << port << std::endl;
        
        // 检测ASIO的实际Windows优化机制
        #ifdef ASIO_HAS_IOCP
            std::cout << "ASIO Windows优化: IOCP (I/O Completion Ports) 已启用" << std::endl;
        #else
            std::cout << "ASIO Windows优化: 使用标准select模式" << std::endl;
        #endif
#else
        std::cout << "UDP 服务器启动，监听端口: " << port << std::endl;
#endif
        
        std::cout << "等待iperf客户端连接..." << std::endl;
        std::cout << "使用命令: iperf -c <服务器IP> -p " << port << " -u -b 1000M -l 1288 -P 1 -t 10" << std::endl;
        std::cout << "按 Ctrl+C 停止服务器" << std::endl;
        std::cout << "----------------------------------------" << std::endl;
        
#ifdef ENABLE_IPERF2_ANALYSIS
        std::cout << "iperf2数据包分析功能: 已启用" << std::endl;
#else
        std::cout << "iperf2数据包分析功能: 已禁用" << std::endl;
#endif
        
        // 显示缓冲区配置信息
        std::cout << "缓冲区配置: ";
#ifdef UDP_BUFFER_CONFIG_STANDARD
        std::cout << "标准配置";
#elif defined(UDP_BUFFER_CONFIG_HIGH_PERFORMANCE)
        std::cout << "高性能配置";
#elif defined(UDP_BUFFER_CONFIG_EXTREME)
        std::cout << "极限性能配置";
#elif defined(UDP_BUFFER_CONFIG_CUSTOM)
        std::cout << "自定义配置";
#else
        std::cout << "默认配置";
#endif
        std::cout << std::endl;
        
        std::cout << "应用层缓冲区: " << UDP_BUFFER_COUNT << " 个 × " 
                  << (UDP_MAX_PACKET_SIZE / 1024) << "KB = " 
                  << (UDP_BUFFER_COUNT * UDP_MAX_PACKET_SIZE / 1024 / 1024) << "MB" << std::endl;
        std::cout << "系统接收缓冲区: " << (UDP_SYSTEM_RCV_BUFFER / 1024 / 1024) << "MB" << std::endl;
        std::cout << "系统发送缓冲区: " << (UDP_SYSTEM_SND_BUFFER / 1024 / 1024) << "MB" << std::endl;
        std::cout << "----------------------------------------" << std::endl;
        
        setup_optimized_buffers();
        setup_signal_handling();
        
        // 启动多个并发接收操作
        for (size_t i = 0; i < buffer_count; ++i) {
            do_receive();
        }
        
        start_statistics_timer();
        
    } catch (const std::exception& e) {
        std::cerr << "socket初始化失败: " << e.what() << std::endl;
        
        // 如果是端口占用问题，给出具体建议
        if (std::string(e.what()).find("bind") != std::string::npos) {
            std::cerr << "端口 " << port << " 可能被占用。请尝试：" << std::endl;
            std::cerr << "1. 关闭其他使用该端口的程序" << std::endl;
            std::cerr << "2. 等待几秒钟后重试" << std::endl;
            std::cerr << "3. 或使用不同的端口" << std::endl;
            std::cerr << "4. 以管理员权限运行程序" << std::endl;
        }
        throw;
    }
}

#ifdef _WIN32
// Windows Console Control Handler实现
BOOL WINAPI UdpReceiver::console_ctrl_handler(DWORD ctrl_type) {
    switch (ctrl_type) {
        case CTRL_C_EVENT:
        case CTRL_BREAK_EVENT:
        case CTRL_CLOSE_EVENT:
        case CTRL_SHUTDOWN_EVENT:
        case CTRL_LOGOFF_EVENT:
            std::cout << "\n收到Windows Control信号，正在安全关闭服务器..." << std::endl;
            
            if (instance_) {
                try {
                    instance_->print_statistics();
                    instance_->stop();
                    instance_->io_context_.stop();
                } catch (const std::exception& e) {
                    std::cerr << "关闭时发生异常: " << e.what() << std::endl;
                }
            }
            
            // 返回TRUE表示我们已处理该信号
            return TRUE;
        default:
            return FALSE;
    }
}
#endif

void UdpReceiver::setup_optimized_buffers() {
    try {
        // 预分配多个缓冲区以提高并发性能
        buffer_pool_.resize(buffer_count);
        
        for (size_t i = 0; i < buffer_count; ++i) {
            buffer_pool_[i].resize(max_length);
        }
        
        // 移除输出信息以保持界面简洁
                  
    } catch (const std::exception& e) {
        std::cerr << "缓冲区设置失败: " << e.what() << std::endl;
        throw;
    }
}

void UdpReceiver::setup_signal_handling() {
    // 注册ASIO信号处理（作为备用）
    signals_.add(SIGINT);
    signals_.add(SIGTERM);
#ifdef SIGQUIT
    signals_.add(SIGQUIT);
#endif

    signals_.async_wait(
        [this](const asio::error_code& error, int signal_number) {
            handle_signal(error, signal_number);
        });
}

void UdpReceiver::handle_signal(const asio::error_code& error, int signal_number) {
    if (!error) {
        std::cout << "\n收到ASIO信号 " << signal_number << "，正在关闭服务器..." << std::endl;
        
        try {
            print_statistics();
            stop();
        } catch (const std::exception& e) {
            std::cerr << "关闭时发生异常: " << e.what() << std::endl;
        }
        
        io_context_.stop();
    }
}

void UdpReceiver::do_receive() {
    if (!running_) return;
    
    // 使用缓冲区池中的缓冲区
    size_t buffer_idx = current_buffer_index_.fetch_add(1) % buffer_count;
    auto& buffer = buffer_pool_[buffer_idx];
    
    // 创建一个shared_ptr来管理异步操作的生命周期
    auto buffer_ptr = std::make_shared<std::vector<char>>(std::move(buffer));
    
    socket_.async_receive_from(
        asio::buffer(*buffer_ptr),
        sender_endpoint_,
        [this, buffer_ptr, buffer_idx](const asio::error_code& ec, std::size_t bytes_recvd) {
            if (!ec && bytes_recvd > 0 && running_) {
                // 记录接收时间
                last_receive_time_ = std::chrono::steady_clock::now();
                
                // 第一次接收数据的处理
                if (!first_packet_received_) {
                    first_packet_received_ = true;
                    test_start_time_ = last_receive_time_;
                    
                    // 显示开始接收数据的时间
                    auto now = std::chrono::system_clock::now();
                    auto time_t = std::chrono::system_clock::to_time_t(now);
                    auto tm = *std::localtime(&time_t);
                    
                    std::cout << "开始接收数据... \"" << std::put_time(&tm, "%H:%M:%S") << "\"" << std::endl;
                }
                
                // 检测新客户端连接
                std::string current_ip = sender_endpoint_.address().to_string();
                unsigned short current_port = sender_endpoint_.port();
                
                if (!client_connected_ || client_ip_ != current_ip || client_port_ != current_port) {
                    client_connected_ = true;
                    client_ip_ = current_ip;
                    client_port_ = current_port;
                    packet_size_ = bytes_recvd;
                    
                    std::cout << "客户端连接: \"" << client_ip_ << "\" : " << client_port_ << std::endl;
                    std::cout << "数据包大小: " << packet_size_ << " 字节" << std::endl;
                    
#ifdef ENABLE_IPERF2_ANALYSIS
                    // 重置序列号统计（新客户端）
                    received_sequences_.clear();
                    min_sequence_ = UINT32_MAX;
                    max_sequence_ = 0;
                    duplicate_packets_ = 0;
                    out_of_order_packets_ = 0;
                    last_sequence_ = 0;
                    sequence_initialized_ = false;
#endif
                }
                
#ifdef ENABLE_IPERF2_ANALYSIS
                // 解析iperf2数据包头部
                IperfHeader header;
                if (parse_iperf_header(buffer_ptr->data(), bytes_recvd, header)) {
                    // 尝试多种解析方式
                    const char* raw_data = buffer_ptr->data();
                    
                    // 只保留最基本的正确解析方式
                    uint32_t seq_correct = ntohl(header.sequence_number);
                    
                    // 使用正确的无符号序列号
                    uint32_t seq = seq_correct;
                    
                    // 检查是否为重复数据包
                    if (received_sequences_.count(seq)) {
                        duplicate_packets_++;
                    } else {
                        received_sequences_.insert(seq);
                        
                        // 更新序列号范围
                        if (seq < min_sequence_) min_sequence_ = seq;
                        if (seq > max_sequence_) max_sequence_ = seq;
                        
                        // // 极简调试：仅显示序列号跳跃
                        // if (sequence_initialized_ && (seq - last_sequence_) > 100) {
                        //     std::cout << "序列号大跳跃: " << last_sequence_ << " -> " << seq << std::endl;
                        // }
                        
                        // 检查是否乱序
                        if (sequence_initialized_ && seq < last_sequence_) {
                            out_of_order_packets_++;
                        }
                        last_sequence_ = seq;
                        sequence_initialized_ = true;
                    }
                }
#endif
                
                // 更新统计信息
                total_bytes_received_ += bytes_recvd;
                total_packets_received_++;
                
                // 将缓冲区归还到池中
                if (buffer_idx < buffer_pool_.size()) {
                    buffer_pool_[buffer_idx] = std::move(*buffer_ptr);
                }
                
                // 立即启动下一次接收以保持高并发
                do_receive();
                
            } else if (running_ && ec != asio::error::operation_aborted) {
                // 错误处理：将缓冲区归还并重试
                if (buffer_idx < buffer_pool_.size()) {
                    buffer_pool_[buffer_idx] = std::move(*buffer_ptr);
                }
                
                // 如果有错误但仍在运行，且不是因为停止操作导致的错误，继续尝试接收
                do_receive();
            }
        });
}

#ifdef ENABLE_IPERF2_ANALYSIS
bool UdpReceiver::parse_iperf_header(const char* data, size_t size, IperfHeader& header) {
    // iperf2头部至少需要12字节
    if (size < 12) {
        return false;
    }
    
    // 解析头部（网络字节序）
    memcpy(&header.sequence_number, data, 4);
    memcpy(&header.timestamp_sec, data + 4, 4);
    memcpy(&header.timestamp_usec, data + 8, 4);
    
    return true;
}

void UdpReceiver::analyze_iperf_packet_format(const char* data, size_t size, size_t packet_count) {
    if (size < 16) return;
    
    std::cout << "\n=== iperf2数据包格式分析 (包 #" << packet_count << ") ===" << std::endl;
    
    // 显示原始字节数据 - 修复格式
    std::cout << "原始数据 (前32字节): ";
    for (size_t i = 0; i < std::min(size, size_t(32)); i++) {
        if (i % 4 == 0) std::cout << " ";
        std::cout << std::hex << std::setw(2) << std::setfill('0') 
                  << (static_cast<unsigned char>(data[i]) & 0xFF);
    }
    std::cout << std::dec << std::endl;
    
    // 更清晰的字节显示
    std::cout << "字节详情:" << std::endl;
    for (size_t i = 0; i < std::min(size, size_t(16)); i += 4) {
        std::cout << "  [" << std::setw(2) << i << "-" << std::setw(2) << (i+3) << "] ";
        for (int j = 0; j < 4 && (i + j) < size; j++) {
            std::cout << std::hex << std::setw(2) << std::setfill('0') 
                      << (static_cast<unsigned char>(data[i + j]) & 0xFF) << " ";
        }
        std::cout << std::dec << std::endl;
    }
    
    std::cout << "============================================\n" << std::endl;
}
#endif

void UdpReceiver::start_statistics_timer() {
    if (!running_) return;
    
    statistics_timer_.expires_after(std::chrono::seconds(2));
    statistics_timer_.async_wait([this](const asio::error_code& ec) {
        if (!ec && running_) {
            print_statistics();
            start_statistics_timer();
        }
    });
}

void UdpReceiver::print_statistics() {
    auto now = std::chrono::steady_clock::now();
    auto total_duration = std::chrono::duration_cast<std::chrono::seconds>(now - start_time_);
    auto interval_duration = std::chrono::duration_cast<std::chrono::milliseconds>(now - last_statistics_time_);
    
    uint64_t current_total_bytes = total_bytes_received_;
    uint64_t current_total_packets = total_packets_received_;
    
    // 检查客户端是否断开连接（超过3秒没有收到数据且之前有连接）
    if (client_connected_ && first_packet_received_) {
        auto time_since_last_receive = std::chrono::duration_cast<std::chrono::seconds>(now - last_receive_time_);
        if (time_since_last_receive.count() > 3) {
            // 客户端断开连接，显示最终统计
            auto test_duration = std::chrono::duration_cast<std::chrono::milliseconds>(last_receive_time_ - test_start_time_);
            double test_seconds = test_duration.count() / 1000.0;
            double final_mbps = (current_total_bytes * 8.0) / (test_seconds * 1000000.0);
            
            std::cout << "========================================" << std::endl;
            std::cout << "客户端断开连接，测试结束" << std::endl;
            std::cout << std::fixed << std::setprecision(2)
                      << "最终统计: 总接收: " << current_total_bytes << " bytes (" 
                      << (current_total_bytes / 1024.0 / 1024.0) << " MB)"
                      << " | 数据包: " << current_total_packets 
                      << " | 测试时长: " << test_seconds << " 秒"
                      << " | 平均速率: " << final_mbps << " Mbps" << std::endl;
            
#ifdef ENABLE_IPERF2_ANALYSIS
            // 计算丢包统计（如果有序列号信息）
            if (min_sequence_ != UINT32_MAX && max_sequence_ != 0 && received_sequences_.size() > 0) {
                uint64_t unique_received = received_sequences_.size();
                uint64_t total_received = current_total_packets;  // 包含重复包的总数
                
                // 简化序列号信息
                std::cout << "序列号范围: " << min_sequence_ << " - " << max_sequence_ 
                          << " (唯一:" << unique_received << ", 重复:" << duplicate_packets_ << ")" << std::endl;
                
                // 修正后的统计：不使用序列号范围计算发送包数
                std::cout << "iperf2接收统计: 接收唯一数据包 " << unique_received 
                          << " | 总接收数据包 " << total_received 
                          << " | 重复数据包 " << duplicate_packets_ 
                          << " | 乱序数据包 " << out_of_order_packets_ << std::endl;
                
                // 注释：无法从接收方准确计算发送包数和丢包率
                std::cout << "注意: 发送包数和丢包率需要参考发送方报告，接收方无法准确计算" << std::endl;
            }
#endif
            
            std::cout << "========================================" << std::endl;
            std::cout << "等待新的客户端连接..." << std::endl;
            std::cout << std::endl;
            
            // 重置状态等待新的客户端
            client_connected_ = false;
            first_packet_received_ = false;
            total_bytes_received_ = 0;
            total_packets_received_ = 0;
            last_bytes_received_ = 0;
            last_packets_received_ = 0;
            
#ifdef ENABLE_IPERF2_ANALYSIS
            // 重置iperf2相关统计
            received_sequences_.clear();
            min_sequence_ = UINT32_MAX;
            max_sequence_ = 0;
            duplicate_packets_ = 0;
            out_of_order_packets_ = 0;
            last_sequence_ = 0;
            sequence_initialized_ = false;
#endif
            
            last_statistics_time_ = now;
            return;
        }
    }
    
    // 如果还没有收到数据，不显示统计信息
    if (!first_packet_received_) {
        last_statistics_time_ = now;
        return;
    }
    
    // 计算间隔内的增量
    uint64_t interval_bytes = current_total_bytes - last_bytes_received_;
    uint64_t interval_packets = current_total_packets - last_packets_received_;
    
    // 计算实时速率（基于间隔）
    double interval_seconds = interval_duration.count() / 1000.0;
    double current_mbps = 0.0;
    double current_pps = 0.0;
    
    if (interval_seconds > 0) {
        current_mbps = (interval_bytes * 8.0) / (interval_seconds * 1000000.0);
        current_pps = interval_packets / interval_seconds;
    }
    
    // 计算总体平均速率（基于测试开始时间）
    auto test_duration = std::chrono::duration_cast<std::chrono::milliseconds>(now - test_start_time_);
    double test_seconds = test_duration.count() / 1000.0;
    double avg_mbps = 0.0;
    
    if (test_seconds > 0) {
        avg_mbps = (current_total_bytes * 8.0) / (test_seconds * 1000000.0);
    }
    
    // 使用新的输出格式
    std::cout << std::fixed << std::setprecision(2)
              << "总接收: " << current_total_bytes << " bytes (" 
              << (current_total_bytes / 1024.0 / 1024.0) << " MB)"
              << " | 数据包: " << current_total_packets 
              << " | 当前速率: " << current_mbps << " Mbps"
              << " | 当前PPS: " << current_pps
              << " | 平均速率: " << avg_mbps << " Mbps" << std::endl;
    
    // 更新上次统计的值
    last_bytes_received_ = current_total_bytes;
    last_packets_received_ = current_total_packets;
    last_statistics_time_ = now;
}

void UdpReceiver::stop() {
    running_ = false;
    
    // 安全地关闭socket和取消定时器
    asio::error_code ec;
    socket_.close(ec);
    statistics_timer_.cancel(ec);
    signals_.cancel(ec);
    
    // 清理缓冲区池
    buffer_pool_.clear();
    
    // 重置状态变量
    first_packet_received_ = false;
    client_connected_ = false;
    client_ip_.clear();
    client_port_ = 0;
    packet_size_ = 0;
    
#ifdef ENABLE_IPERF2_ANALYSIS
    // 重置iperf2相关状态
    received_sequences_.clear();
    min_sequence_ = UINT32_MAX;
    max_sequence_ = 0;
    duplicate_packets_ = 0;
    out_of_order_packets_ = 0;
    last_sequence_ = 0;
    sequence_initialized_ = false;
#endif
    
#ifdef _WIN32
    // 清理静态实例指针
    instance_ = nullptr;
    
    // 移除Console Control Handler
    SetConsoleCtrlHandler(console_ctrl_handler, FALSE);
#endif
}

int main() {
    try {
        // 创建 Asio 的核心 I/O 对象
        asio::io_context io_context;

        // 定义服务器要监听的端口（可根据需要修改）
        unsigned short port = 12345;

        // 创建服务器对象
        UdpReceiver receiver(io_context, port);

        // 运行 I/O 服务，这会阻塞并等待异步操作
        io_context.run();

    } catch (const std::exception& e) {
        std::cerr << "异常: " << e.what() << std::endl;
        return 1;
    }

    std::cout << "服务器已关闭" << std::endl;
    return 0;
}

